package main

import (
    "github.com/tecbot/gorocksdb"
    "path"
    "fmt"
    "sync/atomic"
    "time"
    "os"
)

type RocksdbStore struct {
    tableDef       TableDefine
    partition      Partition
    lastReadTime   atomic.Value
    lastWriteTime  atomic.Value
    lastOpenTime   atomic.Value
    storePath      string
    db             *gorocksdb.DB
    wo             *gorocksdb.WriteOptions
    ro             *gorocksdb.ReadOptions
    opt            *gorocksdb.Options
    datcfh, idxcfh *gorocksdb.ColumnFamilyHandle
}

var CFNames = []string{"default", "idx"}

func NewRocksdbStore(tableDef TableDefine, partition Partition, basePath string) (*RocksdbStore, error) {
    if tableDef.DtLength != len(partition.dt) {
        return nil, fmt.Errorf("dt length not match")
    }
    
    opt := gorocksdb.NewDefaultOptions()
    rateLimiter := gorocksdb.NewRateLimiter(1024, 100*1000, 10)
    opt.SetRateLimiter(rateLimiter)
    opt.SetCreateIfMissing(true)
    opt.SetCreateIfMissingColumnFamilies(true)
    
    ro := gorocksdb.NewDefaultReadOptions()
    wo := gorocksdb.NewDefaultWriteOptions()
    return &RocksdbStore{
        wo:        wo,
        ro:        ro,
        opt:       opt,
        tableDef:  tableDef,
        partition: partition,
        storePath: path.Join(basePath, tableDef.Path(), partition.StoreName()),
    }, nil
}

func (tbs *RocksdbStore) GetTableDef() TableDefine {
    return tbs.tableDef
}
func (tbs *RocksdbStore) GetPartition() Partition {
    return tbs.partition
}

func (tbs *RocksdbStore) Write(rows []*Row) error {
    wb := gorocksdb.NewWriteBatch()
    defer wb.Destroy()
    for _, row := range rows {
        wb.PutCF(tbs.datcfh, []byte(row.Key), []byte(row.Value))
        indexes, err := tbs.tableDef.BuildIndex(row)
        if err != nil {
            return err
        }
        for _, idx := range indexes {
            wb.PutCF(tbs.datcfh, []byte(idx), []byte{})
        }
        b, err := row.MarshallIndex()
        if err != nil {
            return err
        }
        wb.PutCF(tbs.idxcfh, []byte(row.Key), b)
    }
    if err := tbs.db.Write(tbs.wo, wb); err != nil {
        return err
    }
    tbs.lastWriteTime.Store(time.Now().Unix())
    return nil
}

func (tbs *RocksdbStore) LastActiveTime() (int64, int64, int64) {
    o := tbs.lastOpenTime.Load().(int64)
    r := tbs.lastReadTime.Load().(int64)
    w := tbs.lastWriteTime.Load().(int64)
    return o, r, w
}

//TODO add test code
func (tbs *RocksdbStore) Delete(k string) error {
    v, err := tbs.db.GetCF(tbs.ro, tbs.idxcfh, []byte(k))
    defer v.Free()
    if err != nil {
        return err
    }
    r := &Row{Key: k}
    if err := r.UnMarshallIndex(v.Data()); err != nil {
        return err
    }
    idxkeys, err := tbs.tableDef.BuildIndex(r)
    if err != nil {
        return err
    }
    wb := gorocksdb.NewWriteBatch()
    defer wb.Destroy()
    wb.DeleteCF(tbs.datcfh, []byte(k))
    wb.DeleteCF(tbs.idxcfh, []byte(k))
    for _, key := range idxkeys {
        wb.DeleteCF(tbs.datcfh, []byte(key))
    }
    if err := tbs.db.Write(tbs.wo, wb); err != nil {
        return err
    }
    tbs.lastWriteTime.Store(time.Now().Unix())
    return nil
}

func (tbs *RocksdbStore) Query(k string) (string, error) {
    v, err := tbs.db.GetCF(tbs.ro, tbs.datcfh, []byte(k))
    defer v.Free()
    if err != nil {
        return "", err
    }
    tbs.lastReadTime.Store(time.Now().Unix())
    return string(v.Data()), nil
}

func (tbs *RocksdbStore) QueryKeys(condition QueryCondition) ([]string, error) {
    var keys [][]string
    if condition.keyPrefix != "" {
        keys = append(keys, tbs.scan([]byte(condition.keyPrefix), false))
    }
    pres, err := tbs.tableDef.BuildQueryIndexPrefixes(condition)
    if err != nil {
        return nil, err
    }
    for _, p := range pres {
        k := tbs.scan([]byte(p), true)
        keys = append(keys, k)
    }
    tbs.lastReadTime.Store(time.Now().Unix())
    return Keys(keys), nil
}

func (tbs *RocksdbStore) QueryValues(condition QueryCondition) (map[string]string, error) {
    ret := make(map[string]string)
    keys, err := tbs.QueryKeys(condition)
    if err != nil {
        return nil, err
    }
    for _, key := range keys {
        value, err := tbs.Query(key)
        if err != nil {
            return nil, err
        }
        ret[key] = value
    }
    tbs.lastReadTime.Store(time.Now().Unix())
    return ret, nil
}

func (tbs *RocksdbStore) scan(prefix []byte, isIndex bool) (ret []string) {
    iter := tbs.db.NewIteratorCF(tbs.ro, tbs.datcfh)
    defer iter.Close()
    for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
        raw := iter.Key().Data()
        if isIndex {
            if key, err := tbs.tableDef.ExtractKey(raw); err == nil && len(key) > 0 {
                ret = append(ret, string(key))
            }
            continue
        }
        if len(raw) > 0 {
            ret = append(ret, string(raw))
        }
    }
    return ret
}

func (tbs *RocksdbStore) Open() error {
    if _, err := os.Stat(tbs.storePath); err != nil {
        if !os.IsNotExist(err) {
            return err
        }
        err = os.MkdirAll(tbs.storePath, 0644)
        if err != nil {
            return err
        }
        
    }
    //db, err := gorocksdb.OpenDb(tbs.opt, tbs.storePath)
    db, cfh, err := gorocksdb.OpenDbColumnFamilies(tbs.opt, tbs.storePath, CFNames, []*gorocksdb.Options{tbs.opt, tbs.opt})
    if err != nil {
        return err
    }
    tbs.db = db
    tbs.datcfh = cfh[0]
    tbs.idxcfh = cfh[1]
    tbs.lastOpenTime.Store(time.Now().Unix())
    return nil
}

func (tbs *RocksdbStore) Close() error {
    if tbs.db != nil {
        tbs.datcfh.Destroy()
        tbs.idxcfh.Destroy()
        tbs.db.Close()
        tbs.db = nil
        tbs.datcfh = nil
        tbs.idxcfh = nil
    }
    return nil
}

func (tbs *RocksdbStore) IsActive() bool {
    return tbs.db != nil && tbs.datcfh != nil && tbs.idxcfh != nil
}
